package com.su02.multi.examples.mlr;

import com.su02.multi.chainrule.Constant;
import com.su02.multi.chainrule.Function;
import com.su02.multi.chainrule.Variable;

import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.tuple.Pair;

import java.util.ArrayList;
import java.util.List;

public final class MultiLinearRegression {
    private MultiLinearRegression() {}

    /**
     * 前向传播，构造MSE损失函数
     *
     * @param matrix (X | Y)
     * @return 损失函数
     */
    public static Function MSELossFunction(TrainMatrix4MLR matrix) {
        Pair<Integer, Integer> shape = matrix.shape();
        int nRow = shape.getLeft();

        double factor = 1 / (2. * nRow);

        Function function = new Constant(0.);
        for (TrainEntry4MLR entry : matrix) {
            function = function.add(MSELossFunction(entry));
        }

        return function.mul(factor);
    }

    private static Function MSELossFunction(TrainEntry4MLR entry) {
        return MSELossFunction(entry.x(), entry.y());
    }

    private static Function MSELossFunction(List<Double> x, double y) {
        int nx = CollectionUtils.size(x);
        int nTheta = nx + 1;
        List<Function> entries = new ArrayList<>(nTheta);
        int p = 0;
        while (p < nx) {
            // theta_i * x_i
            entries.add(new Variable(p).mul(x.get(p)));
            p ++;
        }
        // theta_n
        entries.add(new Variable(p));

        Function f = new Constant(0.);
        for (Function entry : entries) {
            f = f.add(entry);
        }

        f = (f.sub(new Constant(y))).pow(2.);

        return f;
    }
}